import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

mnist = input_data.read_data_sets('MNIST_data', one_hot=True)

batch_size = 128  # 每次训练的个数
n_inputs = 28  # 每次输入的宽度
n_steps = 28  # 每次输入的长度
lr = 1e-3  # 学习率
train_iters = 100000  # 总训练次数

n_hidden_size = 128  # 隐藏层
n_classes = 10  # 输出

# 定义参数 输入
x = tf.placeholder(tf.float32, [batch_size, n_steps, n_inputs])

# 定义参数 输出
y = tf.placeholder(tf.float32, [None, n_classes])

# 定义参数 weight
weights = {

    'in': tf.Variable(tf.truncated_normal([n_inputs, n_hidden_size])),

    'out': tf.Variable(tf.truncated_normal([n_hidden_size, n_classes]))
}

# 定义参数 biases
biases = {

    'in': tf.Variable(tf.constant(0.1,shape= [n_hidden_size])),

    'out': tf.Variable(tf.constant(0.1,shape = [n_classes]))

}

# 循环神经网络
def RNN(x, weights, biases):
    x = tf.reshape(x, [-1, n_inputs])
    x_in = tf.matmul(x, weights['in']) + biases['in']
    x_in = tf.reshape(x_in, [-1, n_steps, n_hidden_size])

    lstm = tf.nn.rnn_cell.BasicLSTMCell(n_hidden_size, forget_bias=1.0, state_is_tuple=True)

    _init_state = lstm.zero_state(batch_size=batch_size, dtype=tf.float32)

    output, states = tf.nn.dynamic_rnn(lstm, x_in, initial_state=_init_state, time_major=False)

    result = tf.matmul(states[1], weights['out']) + biases['out']

    return result

# 预测的值
pred = RNN(x, weights, biases)

# 定义损失函数
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=pred, labels=y))

# 训练
optimizer = tf.train.AdamOptimizer(lr).minimize(cost)

# 正确预测
correct_pred = tf.equal(tf.argmax(pred,1),tf.argmax(y,1))

# 正确率
accuracy = tf.reduce_mean(tf.cast(correct_pred,tf.float32))

# 初始化
init = tf.global_variables_initializer()

saver = tf.train.Saver()

# 训练
with tf.Session() as sess:
    saver.restore(sess, "./model/model.ckpt")

    step = 0
    while step * batch_size < train_iters:
        print("step :" + str(step))
        train_x, train_y = mnist.test.next_batch(batch_size)
        train_x = train_x.reshape([batch_size, n_steps, n_inputs])
        result = sess.run(pred, feed_dict={x: train_x, y: train_y})
        print("正确值", tf.argmax(train_y, 1).eval())
        print("猜测值", tf.argmax(result, 1).eval())
        print("是否正确", tf.equal(tf.argmax(train_y, 1), tf.argmax(result, 1)).eval())
        step = step + 1